package org.kent.w.base.redis;

import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import javax.annotation.Resource;
import java.util.List;

/**
 * redis 限流
 *
 * @author wu.jigang@geely.com
 * @version 1.0
 * @Date :Created by 2020/4/3.
 */
@Slf4j
public class RedisRateLimiter {

    /**
     * 限流脚本
     * 键值 key
     * 参数 limit(限流值)、expireTime(失效时间，单位：毫秒)
     */
    private static final String REDIS_LIMIT_LUA_SCRIPT = "local key = KEYS[1];\n" +
            "local limit = tonumber(ARGV[1]);\n" +
            "local current = tonumber(redis.call('GET', key) or \"0\");\n" +
            "if current + 1 > limit then\n" +
            "    return 0;\n" +
            "else\n" +
            "    redis.call(\"INCRBY\", key, \"1\");\n" +
            "    redis.call(\"PEXPIRE\", key, ARGV[2]);\n" +
            "    return 1;\n" +
            "end;";

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    /** lua限流脚本 **/
    private RedisScript<Long> redisLimitScript;

    public RedisRateLimiter() {
        redisLimitScript = new DefaultRedisScript<>(REDIS_LIMIT_LUA_SCRIPT, Long.class);
    }

    /**
     * 执行脚本
     *
     * @param key
     * @param limit
     * @param expireTime
     */
    public Long executeScript(String key, Long limit, Long expireTime) {

        List<String> keyList = Lists.newArrayList(key);
        return stringRedisTemplate.execute(redisLimitScript, keyList, limit.toString(), expireTime.toString());
    }

    /**
     * 尝试获取锁
     *
     * @param speed 生成令牌的速率(每秒生成令牌的个数)
     * @param windowTime 令牌时间窗口，令牌累积的时间窗口
     * @param expireTime redis键值的失效时间
     * @return 是否获取到锁
     */
    public Boolean tryAcquire(int speed, int windowTime, long expireTime) {

        // 当前毫秒数
        long currentMills = System.currentTimeMillis();
        // 当前时间窗口已经过的时间
        long currentWindowInterval = currentMills % (windowTime * 1000);
        // 当前时间窗口的起始毫秒数
        long currentWindowStart = currentMills - currentWindowInterval;

        // redis键值：前缀+当前时间窗口起始毫秒数
        String key = "test:rate:limit:" + currentWindowStart;

        // 计算当前时间点应该有的令牌数量，令牌数量 = 当前时间窗口已经过的时间 * speed / 1000
        long limit = currentWindowInterval * speed / 1000;

        log.info("当前毫秒数[{}], 窗口起始毫秒数[{}],窗口已经过的时间[{}], 应生成令牌数[{}]",
                currentMills, currentWindowStart, currentWindowInterval, limit);

        // 执行lua脚本
        Long result = executeScript(key, limit, expireTime);

        return null != result && result == 1;
    }

    /**
     * 调整速率
     *
     * @param speed
     */
    public void adjustDown(int speed) {

        log.warn("暂不支持");
    }
}
